package com.example.watermark;

import org.apache.flink.api.common.eventtime.Watermark;
import org.apache.flink.api.common.eventtime.WatermarkGenerator;
import org.apache.flink.api.common.eventtime.WatermarkGeneratorSupplier;
import org.apache.flink.api.common.eventtime.WatermarkOutput;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.datastream.WindowedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.WindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer;
import org.apache.flink.util.Collector;

import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Properties;

/**
 * @author lbq
 * @version 1.0
 * @date 2022/3/26 21:08
 */
public class WatermarkDemo {
    public static void main(String[] args) throws Exception {
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);
        // env.enableCheckpointing(10);
        env.getConfig().setAutoWatermarkInterval(1000);
        env.setParallelism(3);

        final Properties properties = new Properties();
        properties.setProperty("bootstrap.servers", "aliyun:9092");

        final FlinkKafkaConsumer<String> kafkaConsumer = new FlinkKafkaConsumer<>("flink-demo", new SimpleStringSchema(), properties);

        final DataStreamSource<String> data = env.addSource(kafkaConsumer);

        final SingleOutputStreamOperator<Tuple2<String, Long>> mapStream = data.map(new MapFunction<String, Tuple2<String, Long>>() {
            @Override
            public Tuple2<String, Long> map(String value) throws Exception {
                final String[] fields = value.split(",");
                return new Tuple2<>(fields[0], Long.parseLong(fields[1]));
            }
        });
        final SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS");
        final SingleOutputStreamOperator<Tuple2<String, Long>> watermarks = mapStream.assignTimestampsAndWatermarks(new WatermarkStrategy<Tuple2<String, Long>>() {
            @Override
            public WatermarkGenerator<Tuple2<String, Long>> createWatermarkGenerator(WatermarkGeneratorSupplier.Context context) {
                return new WatermarkGenerator<Tuple2<String, Long>>() {
                    long maxTimeStamp = Long.MIN_VALUE;

                    @Override
                    public void onEvent(Tuple2<String, Long> event, long eventTimestamp, WatermarkOutput output) {
                        maxTimeStamp = Math.max(maxTimeStamp, event.f1);
                        System.out.println("maxTimeStamp: " + maxTimeStamp + "...format:" + sdf.format(maxTimeStamp));
                    }

                    @Override
                    public void onPeriodicEmit(WatermarkOutput output) {
                        System.out.println("...onPeriodicEmit...");
                        output.emitWatermark(new Watermark(maxTimeStamp - 2000));
                    }
                };
            }
        }.withTimestampAssigner((element, recordTimeStamp) -> element.f1));

        final WindowedStream<Tuple2<String, Long>, String, TimeWindow> window = watermarks.keyBy(value -> value.f0).window(TumblingEventTimeWindows.of(Time.seconds(5)));
        window.apply(new WindowFunction<Tuple2<String, Long>, String, String, TimeWindow>() {
            @Override
            public void apply(String s, TimeWindow window, Iterable<Tuple2<String, Long>> input, Collector<String> out) throws Exception {
                final Iterator<Tuple2<String, Long>> iterator = input.iterator();
                List<Long> list = new ArrayList<>();
                while (iterator.hasNext()) {
                    final Tuple2<String, Long> next = iterator.next();
                    list.add(next.f1);
                }
                Collections.sort(list);
                String result = "key:" + s + "..." + "list.size:" + list.size() + "..." + "list.first:" + sdf.format(list.get(0)) + "list.last:" + sdf.format(list.get(list.size() - 1))
                        + "...window.start:" + sdf.format(window.getStart()) + "...window.end:" + sdf.format(window.getEnd());
                out.collect(result);
            }
        }).print();

        env.execute();

    }
}
